package com.zhinengshidai.mcp.servers.command;

import com.alibaba.fastjson.JSONObject;
import com.zhinengshidai.mcp.servers.common.McpServerManager;
import com.zhinengshidai.mcp.servers.common.annotation.Tool;
import io.modelcontextprotocol.spec.McpSchema;

import java.io.BufferedReader;
import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.TimeUnit;

/**
 * @Description
 * @Author chenss
 * @CreateTime 2025-03-24 15:21:17
 * @ModifyTime
 */
public class CommandMCPServer {

    private final static String commandSchema = """
            {
              "type" : "object",
              "properties" : {
                "command" : {
                  "type" : "String"
                },
                "envs" : {
                  "type" : "object"
                },
                "workHome" : {
                  "type" : "string"
                }
              }
            }
            """;

    public static String commandType;

    static {
        if(isWindows()){
            commandType="windows";
        }else {
            commandType="linux";
        }
    }

    public String readResult(BufferedReader bufferedReader) throws IOException {
        StringBuilder stringBuffer = new StringBuilder();
        String result = "";
        while ((result = bufferedReader.readLine()) != null) {
            stringBuffer.append(result);
            stringBuffer.append(",");
        }
        return stringBuffer.toString();
    }

    private static boolean isWindows() {
        String osName = System.getProperty("os.name").toLowerCase();
        return osName.contains("windows");
    }

    private static List<String> buildCommandList(String command) {
        List<String> commandList = new ArrayList<>();
        if (isWindows()) {
            commandList.add("cmd");
            commandList.add("/c");
            commandList.addAll(List.of(command.split(" ")));
        } else {
            commandList.add("sh");
            commandList.add("-c");
            commandList.addAll(List.of(command.split(" ")));
        }
        return commandList;
    }

    @Tool(description = "exec command", schema = commandSchema)
    public String execCommand(Map<String, Object> params) {
        Object envsObj = params.get("envs");
        Map<String, String> envs = new HashMap<>();
        if (Objects.nonNull(envsObj)) {
            envs.putAll((Map<String, String>) envs);
        }
        JSONObject jsonObject = new JSONObject();
        Process process = null;
        try {
            ProcessBuilder processBuilder = new ProcessBuilder();
            processBuilder.environment().putAll(System.getenv());
            processBuilder.environment().putAll(envs);
            processBuilder.redirectErrorStream(true);
            if (Objects.nonNull(params.get("workHome"))) {
                processBuilder.directory(new File(params.get("workHome").toString()));
            }
            processBuilder.command(buildCommandList(params.get("command").toString()));
            process = processBuilder.start();
            String result = readResult(process.inputReader());

            boolean b = process.waitFor(5, TimeUnit.MINUTES);
            if (!b) {
                Thread.sleep(10000);
                if (!process.isAlive()) {
                    if (process.exitValue() == 0) {
                        jsonObject.put("isError", false);
                        jsonObject.put("result", "command exec success");
                        if(!result.isBlank()){
                            jsonObject.put("output", result);
                        }

                    } else {
                        jsonObject.put("isError", true);
                        jsonObject.put("result", "command exec failed");
                        if(!result.isBlank()){
                            jsonObject.put("output", result);
                        }
                    }
                } else {
                    jsonObject.put("isError", false);
                    jsonObject.put("result", "command exec success");
                    if(!result.isBlank()){
                        jsonObject.put("output", result);
                    }
                }
            } else {
                if (process.exitValue() == 0) {
                    jsonObject.put("isError", false);
                    jsonObject.put("result", "command exec success");
                    if(!result.isBlank()){
                        jsonObject.put("output", result);
                    }
                } else {
                    jsonObject.put("isError", true);
                    jsonObject.put("result", "command exec failed");
                    if(!result.isBlank()){
                        jsonObject.put("output", result);
                    }
                }
            }
        } catch (Exception e) {
            jsonObject.put("isError", true);
            jsonObject.put("result", "command exec failed");
            jsonObject.put("output", JSONObject.toJSONString(e));

        }
        return jsonObject.toString();
    }

    public static void main(String[] args) throws Exception {
        McpServerManager mcpServerManager = new McpServerManager();
        mcpServerManager.initStdioServer("server-command", "1.0.0",
                CommandMCPServer.class, McpSchema.ServerCapabilities.builder()
                .tools(true)
                .build());
    }
}
